import numpy as np
from .reward_function_base import BaseRewardFunction


class RelativeAltitudeReward(BaseRewardFunction):
    """
    RelativeAltitudeReward
    Punish if current fighter doesn't satisfy some constraints. Typically negative.
    - Punishment of relative altitude when larger than 1000  (range: [-1, 0])
    
    RelativeAltitudeReward 类是一个奖励函数的子类,它继承自 BaseRewardFunction 基类。该奖励函数用于计算代理的相对高度,并根据高度是否满足一定条件来惩罚代理。以下是该类的详细分析:

    __init__ 方法:在这个方法中,子类调用了基类的构造函数 super().__init__(config) 来初始化基类的属性。子类还初始化了一个额外的属性 KH,该属性用于指定高度惩罚的阈值。如果代理的相对高度超过这个阈值,将会受到惩罚。

    get_reward 方法:这个方法用于计算奖励值,具体步骤如下:
        提取了自己代理的相对高度 ego_z,并将其单位转换为千米(km)。
        提取了敌方代理的相对高度 enm_z,同样将其单位转换为千米(km)。
        计算了相对高度的差值 ego_z - enm_z,表示自己代理与敌方代理之间的高度差。
        计算了奖励值 new_reward,奖励值为 KH - |ego_z - enm_z| 的最小值,即奖励值会在 KH 和 |ego_z - enm_z| 中取最小值,用于惩罚代理当其相对高度超过 KH 时。

    返回奖励值 new_reward。

    总结:RelativeAltitudeReward 类是一个用于计算相对高度奖励的奖励函数。当代理的相对高度超过指定的阈值 KH 时,将会受到惩罚,否则奖励值为零。这个奖励函数的设计可以用来鼓励代理维持与敌方代理之间的相对高度在一定范围内。

    NOTE:
    - Only support one-to-one environments.
    """
    def __init__(self, config):
        super().__init__(config)
        self.KH = getattr(self.config, f'{self.__class__.__name__}_KH', 1.0)     # km

    def get_reward(self, task, env, agent_id):
        """
        Reward is the sum of all the punishments.

        Args:
            task: task instance
            env: environment instance

        Returns:
            (float): reward
        """
        ego_z = env.agents[agent_id].get_position()[-1] / 1000    # unit: km
        enm_z = env.agents[agent_id].enemies[0].get_position()[-1] / 1000    # unit: km
        new_reward = min(self.KH - np.abs(ego_z - enm_z), 0)
        return self._process(new_reward, agent_id)
